import pulp
import numpy as np

def optimize_partition(arr, num_groups):
    prob = pulp.LpProblem("Array Partitioning", pulp.LpMinimize)

    n = len(arr)
    groups = range(num_groups)

    # 决策变量
    x = pulp.LpVariable.dicts("group",
                              ((i, j) for i in range(n) for j in groups),
                              cat='Binary')

    # 组和变量
    group_sums = [pulp.LpVariable(f"group_sum_{j}", lowBound=0) for j in groups]

    # 偏差变量（用于和）
    sum_deviations = [pulp.LpVariable(f"sum_deviation_{j}", lowBound=0) for j in groups]

    # 计数变量（用于记录每组的元素数量）
    group_counts = [pulp.LpVariable(f"group_count_{j}", lowBound=0, cat='Integer') for j in groups]

    # 偏差变量（用于数量）
    count_deviations = [pulp.LpVariable(f"count_deviation_{j}", lowBound=0) for j in groups]

    # 目标值（平均值）
    target_sum = sum(arr) / num_groups
    target_count = n / num_groups

    # 添加约束条件
    for j in groups:
        # 每组的和
        prob += group_sums[j] == pulp.lpSum([arr[i] * x[i, j] for i in range(n)])
        # 和的偏差
        prob += group_sums[j] - target_sum <= sum_deviations[j]
        prob += target_sum - group_sums[j] <= sum_deviations[j]
        # 每组的元素数量
        prob += group_counts[j] == pulp.lpSum([x[i, j] for i in range(n)])
        # 数量的偏差
        prob += group_counts[j] - target_count <= count_deviations[j]
        prob += target_count - group_counts[j] <= count_deviations[j]

    # 每个元素必须属于一个组
    for i in range(n):
        prob += pulp.lpSum(x[i, j] for j in groups) == 1

    # 最小化总和偏差与数量偏差的加权和
    # 可以调整权重 w 来平衡两个目标的重要性
    w = 0.5  # 调整这个参数可以改变对两种偏差的重视程度
    prob += w * pulp.lpSum(sum_deviations) + (1 - w) * pulp.lpSum(count_deviations)

    # 设置求解器并求解
    solver = pulp.PULP_CBC_CMD(msg=1, timeLimit=10)
    prob.solve(solver)

    if pulp.LpStatus[prob.status] != 'Optimal':
        print("未能找到最优解")
        return None

    # 提取结果
    result_groups = [[] for _ in groups]
    for i in range(n):
        for j in groups:
            if pulp.value(x[i, j]) > 0.5:
                result_groups[j].append(arr[i])

    return result_groups


# 输入数组
arr = [12.87, 12.87, 12.87, 13.59, 17.01, 39.87, 43.74, 43.74, 44.46, 44.46, 47.07, 51.39, 51.39, 67.14, 111.87, 39.87,
       43.74, 43.74, 44.01, 44.46, 47.07, 50.4, 51.39, 67.14, 111.87, 114.39]

# 分组数量
num_groups = 3  # 这里可以修改为所需的分组数量

# 优化分组
result = optimize_partition(arr, num_groups)

if result is not None:
    # 打印结果
    for i, group in enumerate(result):
        print(f"组 {i + 1}: {group}")
        print(f"组 {i + 1} 和: {sum(group)}")
        print(f"组 {i + 1} 元素数量: {len(group)}")

    print(f"\n总和: {sum(arr)}")
    print(f"方差: {np.var([sum(group) for group in result])}")
    print(f"元素数量方差: {np.var([len(group) for group in result])}")